# 此脚本为tvm环境

import onnx   # 1.6.0
import numpy as np
import tvm
import tvm.relay as relay
# from tvm.contrib.download import download_testdata

onnx_model = onnx.load('inception-v4.onnx')
from PIL import Image

img_path = 'cat.png'
img = Image.open(img_path).resize((299, 299))
img_ycbcr = img.convert("YCbCr")  # convert to YCbCr
img_y, img_cb, img_cr = img_ycbcr.split()
x = np.array(img_y)[np.newaxis, np.newaxis, :, :]

target = 'cuda'
input_name = 'input'
shape_dict = {input_name: x.shape}
sym, params = relay.frontend.from_onnx(onnx_model, shape_dict)


with relay.build_config(opt_level=3):
    intrp = relay.build_module.create_executor('graph', sym, tvm.gpu(0), target)

dtype = 'float32'
tvm_output = intrp.evaluate(sym)(tvm.nd.array(x.astype(dtype)), **params).asnumpy()
print(tvm_output)